import os
from dotenv import load_dotenv
from db.models.llm_profile.crud import get_llm_model_by_id
from llms.llms_ev.http.llm_chat import ChatZhipuAI
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.output_parsers import JsonOutputParser

# 网页搜索
from langchain_community.utilities import SerpAPIWrapper
# 向量数据库
from langchain_community.vectorstores import Qdrant
from qdrant_client import QdrantClient
from llms.llms_ev.http.embeddings import TaiChiEmbeddings
from langchain.agents import Tool, tool
import requests

# 加载环境变量
load_dotenv()

# 获取SerpAPI的API Key
serper_key = os.getenv("SERPER_KEY")

# 获取API Key
api_key = get_llm_model_by_id(llm_model_id=1)["api_key"]

if not serper_key:
    raise ValueError("SERPER_KEY is not set in environment variables")


@tool  # tool装饰器将函数转换为LangChain工具
def searchs(query: str):
    """只有需要了解实时信息或不知道的事情的时候才会使用这个工具。"""
    if not query:
        return "错误：搜索查询不能为空。请提供具体的查询内容。"

    print(f"执行搜索，查询内容：{query}")
    try:
        serper_tool = SerpAPIWrapper(serpapi_api_key=serper_key)
        res = serper_tool.run(query=query)
        print({"实时搜索信息": res})
        return res
    except Exception as e:
        error_msg = f"搜索执行错误：{str(e)}"
        print(error_msg)
        return error_msg


# 获取Qdrant的配置信息
qdrant_host = os.getenv("QDRANT_ENDPOINT")
qdrant_port = os.getenv("QDRANT_PORT")
qdrant_api_key = os.getenv("QDRANT_API_KEY")


@tool  # tool装饰器将函数转换为LangChain工具
def get_info_from_vdb(query: str):
    """只有回答与2024年运势或者龙年运势相关的问题的时候，会使用这个工具，必须输入用户的生日."""
    client = Qdrant(
        QdrantClient(
            # host=qdrant_host,
            # port=qdrant_port,
            path=os.path.join(os.getcwd(), "../local_qdrant"),
            api_key=qdrant_api_key),
        "local_documents",
        TaiChiEmbeddings(api_key=api_key),
    )
    retriever = client.as_retriever(search_type="mmr", k=3)
    res = retriever.get_relevant_documents(query)
    return res


@tool
def cesuan_api(query: str):
    """只有做八字排盘的时候才会使用这个工具, 需要输入用户姓名和出生年月日时，如果缺少用户姓名和出生年月日时则不可用."""
    url = 'https://api.yuanfenju.com/index.php/v1/Bazi/cesuan'

    # 正确的模板定义
    prompt = ChatPromptTemplate.from_template(
        """你是一个参数查询助手，根据用户输入内容找出相关的参数并按json格式返回。JSON字段如下： 
        - "api_key":"K0I5WCmce7jlMZzTw7vi1xsn0", 
        - "name":"姓名", 
        - "sex":"性别，0表示男，1表示女，根据姓名判断", 
        - "type":"日历类型，0农历，1公历，默认1", 
        - "year":"出生年份 例：1998", 
        - "month":"出生月份 例 8", 
        - "day":"出生日期，例：8", 
        - "hours":"出生小时 例 14", 
        - "minute":"默认就是0" 
        如果没有找到相关参数，则需要提醒用户告诉你这些内容，只返回数据结构，不要有其他的评论。用户输入:{query}"""
    )

    parser = JsonOutputParser()
    prompt = prompt.partial(format_instructions=parser.get_format_instructions())
    chain = prompt | ChatZhipuAI(api_key=api_key) | parser

    try:
        data = chain.invoke({"query": query})
        print("八字查询结果:", data)
    except Exception as e:
        return f"解析错误: {str(e)}"

    try:
        result = requests.post(url, json=data)  # 使用 json 而不是 data 来传递 JSON 数据
        if result.status_code == 200:
            print("====返回数据=====")
            response_json = result.json()
            print(response_json)
            try:
                bazi_info = response_json.get("data", {}).get("bazi_info", {})
                return "八字为:" + bazi_info.get("bazi", "未知")
            except KeyError:
                return "八字查询失败，可能是输入信息不完整。"
        else:
            return "技术错误，请告诉用户稍后再试。"
    except requests.exceptions.RequestException as e:
        return f"API 请求失败: {str(e)}"

# 实例化搜索工具，Tool类是LangChain库中定义的工具类，用于定义一个工具，该工具可以执行特定的任务。
# search_tool = Tool(
#     name="search",
#     func=searchs,
#     description="只有需要了解互联网实时信息时或者不知道的事情才调用"
# )
#
# get_info_from_vdb_tool = Tool(
#     name="get_info_from_vdb",
#     func=get_info_from_vdb,
#     description="当用户希望了解龙年运势的时候，你会查询本地知识库工具"
# )
